import numpy as np
from numpy.linalg import inv, lstsq
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score

# 一元数据准备
train_X1 = np.array([[6],[8],[10],[14],[18]])
train_Y1 = [[7],[9],[13],[17.5],[18]]
test_X1 = np.array([[8],[9],[11],[16],[12]])
test_Y1 = [[11],[8.5],[15],[18],[11]]

# 多元数据准备
train_X2 = [
    [1,6,2],[1,8,1],[1,10,0],[1,14,2],[1,18,0]
] 
train_X2_1 = np.array([[6],[8],[10],[14],[18]])
train_X2_2 = np.array([[2],[1],[0],[2],[0]])
train_Y2 = [[7],[9],[13],[17.5],[18]]
test_X2 = [
    [1,8,2],[1,9,0],[1,11,2],[1,16,2],[1,12,0]
]
test_X2_1 = np.array([[8],[9],[11],[16],[12]])
test_X2_2 = np.array([[2],[0],[2],[2],[0]])
test_Y2 = [[11],[8.5],[15],[18],[11]]



# 一元线性回归
model1 = LinearRegression()
model1.fit(train_X1, train_Y1)
print("一元线性回归的score:", model1.score(test_X1, test_Y1))
# print("一元线性回归的系数:", model1.coef_[0])
# print("一元线性回归的截距:", model1.intercept_)

# 一元非线性回归
X = np.column_stack([np.array(train_X1)**2, train_X1, np.ones_like(train_X1)])
coefficients = lstsq(X, train_Y1,rcond=None)[0]
print("一元非线性回归的系数：",coefficients[0], coefficients[1], coefficients[2])
# 一元非线性回归的评估
y1_pred_train = coefficients[0] * train_X1**2 + coefficients[1] * train_X1 + coefficients[2]
y1_pred_test = coefficients[0] * test_X1**2 + coefficients[1] * test_X1 + coefficients[2]
print("一元非线性回归的评估:", r2_score(train_Y1, y1_pred_train))
print("一元非线性回归的预测:", r2_score(test_Y1, y1_pred_test))

# 多元线性回归
a, b1, b2 = np.dot(inv(np.dot(np.transpose(train_X2), train_X2)), 
          np.dot(np.transpose(train_X2), train_Y2))
print("多元线性回归的系数：", a, b1, b2)

model_2 = LinearRegression()
model_2.fit(train_X2, train_Y2)

model_2.predict(test_X2)
print("多元线性回归的预测:", model_2.score(test_X2, test_Y2))

# 多元非线性回归
X2 = np.column_stack([train_X2_1**2, train_X2_2**2, train_X2_1, train_X2_2, np.ones_like(train_X2_1)])
coefficients2 = lstsq(X2, train_Y2, rcond=None)[0]
print("多元非线性回归的系数", coefficients2)
# 多元非线性回归的评估
y2_pred_train = coefficients2[0] * train_X2_1**2 + coefficients2[1] * train_X2_2**2 + coefficients2[2] * train_X2_1 + coefficients2[3] * train_X2_2 + coefficients2[4]
y2_pred_test = coefficients2[0] * test_X2_1**2 + coefficients2[1] * test_X2_2**2 + coefficients2[2] * test_X2_1 + coefficients2[3] * test_X2_2 + coefficients2[4]
print("多元非线性回归的评估：", r2_score(train_Y2, y2_pred_train))
print("多元非线性回归的预测：", r2_score(test_Y2, y2_pred_test))



# 可视化
plt.scatter(train_X1, train_Y1, color='blue')
plt.plot(test_X1, model1.predict(test_X1), color='red')
plt.savefig('1.png')